import pandas as pd
from pandasai import Agent
from pandasai.llm.openai import OpenAI

# 读取 CSV 文件，使用原始字符串避免转义问题
df = pd.read_csv(r"D:\宁德时代股票.csv")
# 将 日期 列转换为日期时间格式
df['日期'] = pd.to_datetime(df['日期'])
# 设置 Date 列为索引
df.set_index('日期', inplace=True)

# 配置 Deepseek 模型
OpenAI._supported_chat_models.append('deepseek-chat')
llm_openai = OpenAI(
    api_base='https://api.deepseek.com/v1',
    api_token='sk-d57a50bd4d19428583093ee15d3e4510',  # 替换为实际的 API 令牌
    model='deepseek-chat'
)

# 创建 Agent
agent = Agent(dfs=[df], config={"llm": llm_openai})


# 生成可视化图表并标注历史最高点和最低点，假设股票价格列名为 'Price'
visualization_prompt = """
请绘制一张折线图，展示数据框中 'Price' 列随时间（索引）的变化情况，
X 轴为日期，Y 轴为 股票价格，图表标题为 ' stock price trend'。
同时，请在图中标注出 'Price' 列的历史最高点和最低点，
分别用红色三角形和绿色三角形表示，并在旁边添加文本标注具体数值。
"""
agent.chat(visualization_prompt)

